package ar.edu.itba.pod.legajo46470.balance;

import static com.google.common.base.Preconditions.checkNotNull;

import java.rmi.RemoteException;
import java.rmi.server.UnicastRemoteObject;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.log4j.Logger;
import org.joda.time.DateTime;

import ar.edu.itba.balance.api.AgentsBalancer;
import ar.edu.itba.balance.api.AgentsTransfer;
import ar.edu.itba.balance.api.NodeAgent;
import ar.edu.itba.balance.api.NotCoordinatorException;
import ar.edu.itba.node.Node;
import ar.edu.itba.node.NodeInformation;
import ar.edu.itba.node.api.ClusterAdministration;
import ar.edu.itba.pod.legajo46470.cluster.node.ClusterNode;
import ar.edu.itba.pod.legajo46470.communication.ClusterCommunication;
import ar.edu.itba.pod.thread.CleanableThread;

public class AgentsBalancerImpl extends UnicastRemoteObject implements AgentsBalancer {
	private static final long serialVersionUID = -2271516439241928825L;
	private final static Logger LOGGER = Logger.getLogger(AgentsBalancerImpl.class);
	private final static long COORDINATORTIME = 3000;
	private ClusterNode selfNode;
	private CountDownLatch coordinatorLock;
	private volatile NodeInformation coordinator;
	private final List<AgentBalanceInfo> balanceInfo; // after balance
	private final Map<NodeInformation, Integer> actualInfo; // original amount of nodes per agent before balance
	private final Map<NodeInformation, Long> lastElectionMessages;
	private final Map<NodeInformation, Long> lastCoordinatorMessages;
	private final ClusterCommunication<AgentsBalancer> clusterCommBalancer;
	private final ClusterCommunication<AgentsTransfer> clusterCommTransfer;
	private final AtomicBoolean onElection;
	
	public AgentsBalancerImpl(ClusterNode selfNode) throws RemoteException {
		checkNotNull(selfNode, "Self Node cannot be null");
		this.selfNode = selfNode;
		clusterCommBalancer = new ClusterCommunication<AgentsBalancer>(selfNode, Node.AGENTS_BALANCER);
		clusterCommTransfer = new ClusterCommunication<AgentsTransfer>(selfNode, Node.AGENTS_TRANSFER);
		lastElectionMessages = new ConcurrentHashMap<NodeInformation, Long>();
		lastCoordinatorMessages = new ConcurrentHashMap<NodeInformation, Long>();
		onElection = new AtomicBoolean(false);
		coordinatorLock = new CountDownLatch(1);
		balanceInfo = new ArrayList<AgentBalanceInfo>();
		actualInfo = new ConcurrentHashMap<NodeInformation, Integer>();
		
		/*new CleanableThread("msgCleaner") {
			@Override
			public void run() {
				try {
					while(true) {
						Thread.sleep(COORDINATORTIME*6);
						lastElectionMessages.clear();
						lastCoordinatorMessages.clear();
					}
				} catch (InterruptedException e) {
					LOGGER.warn("Cleaner thread interrupted!");
				}
			}
		}.start();*/
	}
	
	@Override
	public void bullyElection(NodeInformation node, long timestamp)
			throws RemoteException {
		final NodeInformation sourceNode = node;
		final long sourceTimeStamp = timestamp;
		
		// checks if this message was already received
		if (!messageRepeated(lastElectionMessages, node, timestamp)) {
			LOGGER.info("new election message from: " + node.id() + " at time: " + timestamp);
			//LOGGER.info("is sending node smaller than me? " + String.valueOf(node.id().compareTo(selfNode.getNodeInfo().id()) < 0));
			if (node.id().compareTo(selfNode.getNodeInfo().id()) <= 0) {
				// source node id is either me or weaker, bully him if weaker and start election myself
				if (onElection.compareAndSet(false, true)) {
					new CleanableThread("startElection") {
						@Override
						public void run() {
							
							AgentsBalancer agentBalancerStub;
							
							try {
								// sends bullyOK message to node unless its me
								
								if (!sourceNode.id().equals(selfNode.getNodeInfo().id())) {
									//LOGGER.info("I: " + selfNode.getNodeInfo().id() + " bullying: " + sourceNode.id());
									//LOGGER.info(selfNode.getNodeInfo().id() + " after bullying, starting new election");
									agentBalancerStub = (AgentsBalancer) clusterCommBalancer.getStub(sourceNode);
									agentBalancerStub.bullyOk(sourceNode);
								}
								
								// broadcasts election to others (except self and sender)
								broadcast(selfNode.getNodeInfo(), DateTime.now().getMillis());
								
								Thread.sleep(COORDINATORTIME);
								
								// if onElection, no bullyOK arrived. I won!
								if (onElection.compareAndSet(true, false)) {
									//LOGGER.info("I WON!: " + selfNode.getNodeInfo().id());
									// send coordinator msg to self to replicate to neighbours
									bullyCoordinator(selfNode.getNodeInfo(), DateTime.now().getMillis());
								}
							} catch (RemoteException e) {
								LOGGER.error("Remote Exception" + e.getMessage());
								clusterCommBalancer.disconnectNode(sourceNode);
							} catch (IllegalStateException e) {
								LOGGER.warn("Node down " + sourceNode.id());
							} catch (InterruptedException e) {
								LOGGER.warn("Interrupted while waiting for bullyOK or win election");
							}
								
						}
					}.start();
				} else {
					if (!sourceNode.equals(selfNode.getNodeInfo())) {
						// sends bullyOK message to node
						//LOGGER.info("I: " + selfNode.getNodeInfo().id() + " bullying: " + sourceNode.id());
						AgentsBalancer agentBalancerStub = (AgentsBalancer) clusterCommBalancer.getStub(sourceNode);
						try {
							agentBalancerStub.bullyOk(sourceNode);
						} catch (RemoteException e) {
							// on error calling bullyOK, node disconnection!
							LOGGER.warn("node down: " + sourceNode.id());
							clusterCommBalancer.disconnectNode(sourceNode);
						}
					}
				}
			} else {
				new CleanableThread("broadcastElection") {
					@Override
					public void run() {
						try {
							// broadcasts election to others (except self and sender)
							broadcast(sourceNode, sourceTimeStamp);
						} catch (RemoteException e) {
							// never happens
						}
					}
				}.start();
				
			}
		}
	}

	@Override
	public void bullyOk(NodeInformation node) throws RemoteException {
		onElection.set(false);
		// TODO interrumpir el thread que espera a proclamarse coordinador
	}

	@Override
	public void bullyCoordinator(NodeInformation node, long timestamp)
			throws RemoteException {
		final NodeInformation sourceNode = node;
		final long sourceTimeStamp = timestamp;
		
		if (!messageRepeated(lastCoordinatorMessages, node, timestamp)) {
			coordinator = node;
			onElection.set(false); // election finished
			coordinatorLock.countDown(); //releases lock
			LOGGER.info("The coordinator is: " + node.id());
			
			new CleanableThread("balanceThread") {
				@Override
				public void run() {
					// if i'm the coordinator, i need to balance the cluster
					if (sourceNode.equals(selfNode.getNodeInfo())) {
						// sincronized to avoid balance interference
						synchronized (AgentsBalancerImpl.class) {
							LOGGER.info("starting balance");
							getNodeAgentsQty(null);
							generateBalanceInfo();
						}
					}
				}
			}.start();
			
			new CleanableThread("broadcastCoordinator") {
				@Override
				public void run() {					
					try {
						// broadcasts coordinator to others (except self and sender)
						broadcastCoordinator(sourceNode,sourceTimeStamp);
					} catch (RemoteException e) {
						// never happens
					}
				}
			}.start();
		}
	}

	@Override
	public void shutdown(List<NodeAgent> agents) throws RemoteException,
			NotCoordinatorException {
		if (!coordinator.equals(selfNode.getNodeInfo())) {
			throw new NotCoordinatorException(coordinator);
		}
		
		// if shutting node did not have any agent, there's not point in balancing
		if (agents.size() > 0) {
			synchronized (AgentsBalancerImpl.class) {
				// add all the agents of the node shutting down to one node and then balance
				NodeInformation offlineNode = agents.get(0).node();
				
				// if i'm shutting down and im alone, just finish without any balance
				if (offlineNode.equals(selfNode.getNodeInfo()) && balanceInfo.size() == 1) {
					return;
				}
				balanceInfo.remove(new AgentBalanceInfo(offlineNode, agents.size()));
				
				Collections.sort(balanceInfo);
				
				AgentBalanceInfo agentBalanceInfo = balanceInfo.get(0);
				
				AgentsTransfer agentsTranfer;
				
				LOGGER.info("Shutdown received from node: " + offlineNode + ", transfering his agents to: " + agentBalanceInfo.getNode());
				
				// update balance info list, updating amount of agents
				agentBalanceInfo.setAgentsQty(agents.size() + agentBalanceInfo.getAgentsQty());
				
				// get that node stub and send agents to him
				agentsTranfer = clusterCommTransfer.getStub(agentBalanceInfo.getNode());
				agentsTranfer.runAgentsOnNode(agents);
				
				// balance!
				getNodeAgentsQty(offlineNode);
				
				// updates real info, outgoing node has no agents anymore because they were transfered
				actualInfo.put(offlineNode, 0);
				//actualInfo.put(agentBalanceInfo.getNode(), agentBalanceInfo.getAgentsQty());
				
				LOGGER.info("Post shutdown balance");
				// perform balance!
				generateBalanceInfo();
			}
		}
	}

	@Override
	public void addAgentToCluster(NodeAgent agent) throws RemoteException,
			NotCoordinatorException {
		AgentsTransfer agentsTransfer;
		List<NodeAgent> singleAgentList = new ArrayList<NodeAgent>();
		
		if (!coordinator.equals(selfNode.getNodeInfo())) {
			throw new NotCoordinatorException(coordinator);
		}
		
		singleAgentList.add(agent);
		
		synchronized (AgentsBalancerImpl.class) {
			LOGGER.info("Node: " + agent.node().id() + " sent me an agent to add to the cluster");
			
			if (balanceInfo.size() == 0) {
				balanceInfo.add(new AgentBalanceInfo(selfNode.getNodeInfo(), 0));
			}
			
			Collections.sort(balanceInfo);
			
			
			AgentBalanceInfo ainfo = balanceInfo.get(0);
			
			if (!ainfo.getNode().equals(selfNode.getNodeInfo())) {
				agentsTransfer = (AgentsTransfer) clusterCommTransfer.getStub(ainfo.getNode());
			} else {
				agentsTransfer = selfNode.getAgentsTransfer();
			}
			
			LOGGER.info("Sending agent to node: " + ainfo.getNode());
			agentsTransfer.runAgentsOnNode(singleAgentList);
			// updates logical amount of agents in node
			ainfo.setAgentsQty(ainfo.getAgentsQty()+1);
			Collections.sort(balanceInfo);
		}
	}
	
	public NodeInformation getCoordinator() {
		boolean electionFinished = false;
		
		try { 
			// validation to avoid returning constantly an old or missing coordinator
			clusterCommBalancer.getStub(coordinator);
		} catch (Exception e) {
			// means my coordinator has disconnected
			LOGGER.warn((coordinator == null? "Coordinator not set," : "Coordinator " + coordinator + " unreachable,") +  " restarting coordinator");
			coordinator = null;
		} finally {
			while(coordinator == null) {
				try {
					electionFinished = coordinatorLock.await(COORDINATORTIME, TimeUnit.MILLISECONDS);

					if (!electionFinished) {
						bullyElection(selfNode.getNodeInfo(), DateTime.now().getMillis());
					} else {
						coordinatorLock = new CountDownLatch(1);
					}
				} catch (InterruptedException e) {
					LOGGER.info("CountDownLatch interupted!");
				} catch (RemoteException e) {
					// never happens
				}
			}
		}
		return coordinator;
	}	
	private boolean messageRepeated(Map<NodeInformation, Long> messages, 
			NodeInformation node, long timestamp) {
		synchronized (AgentsBalancerImpl.class) {
			if (!messages.containsKey(node) || messages.get(node).longValue() < timestamp) {
				messages.put(node, DateTime.now().getMillis());
				return false;
			} else {
				return true;
			}
		}
	}
	
	private void getNodeAgentsQty(NodeInformation excludeNode) {
		AgentsTransfer agentsTransfer;
		
		// resets all info related to amount of agents per node
		balanceInfo.clear();
		actualInfo.clear();
		try {
			// ask for my number of agents ONLY if im not the one disconnecting
			if (excludeNode == null || !excludeNode.equals(selfNode.getNodeInfo())) {
				balanceInfo.add(new AgentBalanceInfo(selfNode.getNodeInfo(), selfNode.getAgentsTransfer().getNumberOfAgents()));
			}
			
			for (NodeInformation n : selfNode.getClusterAdministration().connectedNodes()) {
				if (!n.equals(selfNode.getNodeInfo())) {
					if (excludeNode == null || !n.equals(excludeNode)) {
						try {
							// gets all neighbours agent amount and sort list
							agentsTransfer = (AgentsTransfer) clusterCommTransfer.getStub(n);
							balanceInfo.add(new AgentBalanceInfo(n, agentsTransfer.getNumberOfAgents()));
							
		
						} catch (RemoteException e) {
							LOGGER.warn("Node down: " + n.id());
							clusterCommTransfer.disconnectNode(n);
						} catch (IllegalStateException e) {
							LOGGER.warn("Node down: " + n.id());
							//clusterCommTransfer.disconnectNode(n);
						}
					}
				}
			}
			// sort info and store a backup for later use in actualInfo
			Collections.sort(balanceInfo);
			
			for (AgentBalanceInfo a : balanceInfo) {
				actualInfo.put(a.getNode(), a.getAgentsQty());				
			}
		} catch (RemoteException e) {
			// never happens, local call
		}
		
	}
	
	private void generateBalanceInfo() {
		boolean condition;
		int min, max, newAgentsForMax, newAgentsForMin, diff;
		
		// sorts ascending by amount of agents
        Collections.sort(balanceInfo);

        max = balanceInfo.get(balanceInfo.size()-1).getAgentsQty();
        min = balanceInfo.get(0).getAgentsQty();
        diff = max-min;
        
        condition = diff > 1;
        
        // while !balanced
        while(condition){
            newAgentsForMax = max-diff/2;
            newAgentsForMin = min+diff/2;
            balanceInfo.get(balanceInfo.size()-1).setAgentsQty(newAgentsForMax);
            balanceInfo.get(0).setAgentsQty(newAgentsForMin);
            
            // 
            Collections.sort(balanceInfo);
            max = balanceInfo.get(balanceInfo.size()-1).getAgentsQty();
            min = balanceInfo.get(0).getAgentsQty();
            diff = max-min;
            
            condition = diff > 1;
        }
        
        // final ascending sort
        Collections.sort(balanceInfo);
        
        // after calculations, perform real balance
        doBalance();
	}
	
	private void doBalance() {
		int originalAgentAmount, diff;
		List<AgentBalanceInfo> balanceInfoCopy = new ArrayList<AgentBalanceInfo>(balanceInfo);
		List<NodeAgent> bagOfExtractedAgents = new ArrayList<NodeAgent>();
		List<NodeAgent> sublistOfBag = new ArrayList<NodeAgent>();
		AgentsTransfer agentsTransfer;
		
		// extract agents from nodes based on calculations on balanceInfo data
		for (AgentBalanceInfo ainfo : balanceInfoCopy) {
			originalAgentAmount = actualInfo.get(ainfo.getNode());
			diff = originalAgentAmount - ainfo.getAgentsQty();
			
			if (diff > 0) {
				try { 
					agentsTransfer = (AgentsTransfer) clusterCommTransfer.getStub(ainfo.getNode());
					bagOfExtractedAgents.addAll(agentsTransfer.stopAndGet(diff));
					//ainfo.setAgentsQty(ainfo.getAgentsQty() - diff);
				} catch (IllegalStateException e) { 
					balanceInfo.remove(ainfo);
					LOGGER.warn("Node down: " + ainfo.getNode() + ", node agents lost!");
					// TODO por ahora nada, ver dps tolerancia a fallas
					// recordar que ya el getStub desconecto al nodo
				} catch (RemoteException e) {
					balanceInfo.remove(ainfo);
					LOGGER.warn("Node down: " + ainfo.getNode() + ", node agents lost!");
					clusterCommBalancer.disconnectNode(ainfo.getNode());
				}
				
			}
		}
		
		balanceInfoCopy.clear();
		balanceInfoCopy.addAll(balanceInfo);
		
		// assign extra agents to nodes that need them acoording to balanceInfo
		for (AgentBalanceInfo ainfo : balanceInfoCopy) {
			originalAgentAmount = actualInfo.get(ainfo.getNode());
			diff = originalAgentAmount - ainfo.getAgentsQty();
			if (diff < 0 && bagOfExtractedAgents.size() > 0) {
				try {
					
					// sent number of agents dictated by balanceInfo to node, to balance him
					// and remove them from the bag of agents previously extracted
					agentsTransfer = (AgentsTransfer) clusterCommTransfer.getStub(ainfo.getNode());
					//ainfo.setAgentsQty(ainfo.getAgentsQty() + Math.abs(diff));
					
					sublistOfBag.addAll(bagOfExtractedAgents.subList(0, Math.abs(diff)));
					agentsTransfer.runAgentsOnNode(sublistOfBag);
					bagOfExtractedAgents.removeAll(sublistOfBag);
					sublistOfBag.clear();
				} catch (IllegalStateException e) {
					LOGGER.warn("Node down: " + ainfo.getNode());
					balanceInfo.remove(ainfo.getNode());
				} catch (RemoteException e) {
					LOGGER.warn("Node down: " + ainfo.getNode());
					clusterCommBalancer.disconnectNode(ainfo.getNode());
					balanceInfo.remove(ainfo.getNode());
				}
			}
		}
		// check to see if not all agents were delivered due to node disconnections
		checkRemainingInBag(bagOfExtractedAgents, balanceInfoCopy);	
		balanceInfoCopy.clear(); // final reset to release memory usage
		LOGGER.info("Balance finished");
	}	
	
	private void checkRemainingInBag(List<NodeAgent> bagOfExtractedAgents, 
			List<AgentBalanceInfo> balanceInfoCopy) {
		AgentsTransfer agentsTransfer;
	
		balanceInfoCopy.clear(); // resets copy
		balanceInfoCopy.addAll(balanceInfo); // reloads it
		
		// check if remaining agents in bag due to any node disconnection after "logical" balance done
		if (bagOfExtractedAgents.size() > 0) {
			boolean agentsSent = false;
			AgentBalanceInfo ainfo = null;
			
			// assign all of them to myself, rebalance would be bad performance
			LOGGER.warn("extra agents due to node disconnections after balance info generation");
			
			// at least one node will receive the agents (myself)
			for (int i = 0 ; i < balanceInfoCopy.size() && !agentsSent ; i++) {
				try {
					// get the transfer stub, send the nodes and update logical amount of agents in node
					ainfo = balanceInfo.get(i);
					agentsTransfer = (AgentsTransfer) clusterCommTransfer.getStub(ainfo.getNode());
					agentsTransfer.runAgentsOnNode(bagOfExtractedAgents);
					ainfo.setAgentsQty(ainfo.getAgentsQty() + bagOfExtractedAgents.size());
					agentsSent = true;
				} catch (IllegalStateException e) {
					LOGGER.warn("Node down: " + ainfo.getNode());
					balanceInfo.remove(ainfo);
					// node already disconnected
				} catch (RemoteException e) {
					LOGGER.warn("Node down: " + ainfo.getNode());
					clusterCommBalancer.disconnectNode(ainfo.getNode());
				}
			}
		}
	}

	private void broadcast(NodeInformation sourceNode, long sourceTimeStamp) throws RemoteException {
		ClusterAdministration clusterAdministration = selfNode.getClusterAdministration();
		AgentsBalancer agentBalancerStub;
		
		// broadcasts election to others (except self and sender)
		for (NodeInformation n : clusterAdministration.connectedNodes()) {
			try {
				if (!((n.equals(selfNode.getNodeInfo()) || n.equals(sourceNode)))) {
					agentBalancerStub = (AgentsBalancer) clusterCommBalancer.getStub(n);
					agentBalancerStub.bullyElection(sourceNode, sourceTimeStamp);
				}
			} catch (RemoteException e) {
				LOGGER.error("broadcastElection: Unreacheable node: " + n.id());
				clusterCommBalancer.disconnectNode(n);
			} catch (IllegalStateException e) {
				LOGGER.warn("Node down " + n.id());
			}
			
		}
	}
	
	private void broadcastCoordinator(NodeInformation sourceNode, long sourceTimeStamp) throws RemoteException {
		ClusterAdministration clusterAdministration = selfNode.getClusterAdministration();
		AgentsBalancer agentBalancerStub;
		
		for (NodeInformation n : clusterAdministration.connectedNodes()) {
			try {
				if (!((n.equals(selfNode.getNodeInfo()) || n.equals(sourceNode)))) {
					agentBalancerStub = (AgentsBalancer) clusterCommBalancer.getStub(n);
					agentBalancerStub.bullyCoordinator(sourceNode, sourceTimeStamp);
					LOGGER.info("Coordinador sent to: " + n.id());
				}
			} catch (RemoteException e) {
				LOGGER.error("broadcastCoordinator: Unreacheable node: " + n.id());
				clusterCommBalancer.disconnectNode(n);
			} catch (IllegalStateException e) {
				LOGGER.warn("Node down " + n.id());
			}
			
		}
	}
}
